from model.multimodel_regress import model
from tools.infer import infer, infer_single
import torch
from loader.loader import get_loader
import colorama
from sklearn.metrics import roc_auc_score
import os
import numpy as np
from tools.train_regress import evaluate_statistical_coefficient
import csv


def list_files_in_directory(directory):
    try:
        files = os.listdir(directory)
        file_names = [file for file in files if os.path.isfile(os.path.join(directory, file))]
        return file_names
    except Exception as e:
        return str(e)


colorama.init(autoreset=True)
device = torch.device('cuda')
fold_index = 1
model.to(device)
model.eval()
# 定义文件夹路径
results_directory = 'results'
csv_file_path = os.path.join(results_directory, 'regression_5_results.csv')

# 检查文件夹是否存在，如果不存在则创建
if not os.path.exists(results_directory):
    os.makedirs(results_directory)

# 打开 CSV 文件以写入结果
with open(csv_file_path, mode='w', newline='') as csv_file:
    fieldnames = ['Fold', 'File', 'RMSE', 'MAE', 'R2', 'Pearson']
    writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
    writer.writeheader()  # 写入表头

    for fold_index in range(1, 6):
        directory = f'output5/fold_{fold_index}_regression/'
        dir_files = list_files_in_directory(directory)
        print(colorama.Fore.GREEN + f'fold {fold_index}')

        for file in dir_files:
            model.load_state_dict(torch.load(f'{directory}{file}'))
            _, _, test_loader = get_loader(True, fold_index)
            predictions_list = []
            targets_list = []

            with torch.no_grad():
                for images, labels in test_loader:
                    images, labels = images.to(device), labels.to(device)
                    outputs = model(images)
                    predictions_list.extend(outputs.detach().cpu().numpy())
                    targets_list.extend(labels.detach().cpu().numpy())

            predictions = np.array(predictions_list)
            targets = np.array(targets_list)
            rmse, mae, r2, correlation_coefficient, pearson = evaluate_statistical_coefficient(predictions, targets)

            # 打印结果到控制台
            print(colorama.Fore.RED + f'{file} test RMSE: {rmse:.4f}', end='\t')
            print(colorama.Fore.CYAN + f'MAE: {mae:.4f}', end='\t')
            print(colorama.Fore.GREEN + f'R2: {r2:.4f}', end='\t')
            print(colorama.Fore.LIGHTYELLOW_EX + f'Pearson: {correlation_coefficient:.4f}')

            # 将结果写入 CSV 文件
            writer.writerow({
                'Fold': fold_index,
                'File': file,
                'RMSE': rmse,
                'MAE': mae,
                'R2': r2,
                'Pearson': correlation_coefficient
            })
